import os 
import cv2 as cv
import numpy as np
from sklearn.model_selection import train_test_split
from keras.utils.np_utils import to_categorical

if __name__ == "__main__":

    all_data = []
    all_label = []

    for i in range(10):
        path = './picture/%d'%i
        for f in os.listdir(path):

            img = cv.imread(os.path.join(path, f))          #读入图像文件 数组结构 [横坐标(1...n),纵坐标(1...n),颜色(BGR)]
            img = cv.resize(img, (32,32))[...,(2,1,0)]      # 缩放，通道更换
            all_data.append(img)                            #数据追加
            all_label.append(i)                             #标签追加

    all_data = np.asarray(all_data)     #list转化为ndarray
    all_label = np.asarray(all_label)

    data_train, data_test, label_train, label_test = train_test_split(all_data, all_label, test_size=0.2, random_state=40)  # 使用sklearn中的API划分数据集

    data_train = data_train / 128.0 - 1
    label_train = to_categorical(label_train)
    data_test = data_test / 128.0 - 1           #颜色数据（0-255）-->（-1，1）归一化
    label_test = to_categorical(label_test)     #向量化(1*10),配合下文的多分类交叉熵损失函数 0->[1,0,0,0,0,0,0,0,0,0]

    np.save("data_train", data_train)           #保存为npy文件
    np.save("label_train", label_train)
    np.save("data_test", data_test)
    np.save("label_test", label_test)